package com.base.commons.mybatis.sql;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.mysql.jdbc.Connection;
import com.mysql.jdbc.ResultSetInternalMethods;
import com.mysql.jdbc.Statement;
import com.mysql.jdbc.StatementInterceptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;
import java.util.Properties;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Pattern;

/**
 * 输出在 mysql 中运行的完整 sql(mybatis hibernate 都使用的 ? 占位符)
 *
 * 连接参数: &statementInterceptors=com.zt.fbg.commons.mybatis.sql.ShowSqlInterceptor
 */
public class ShowSqlInterceptor implements StatementInterceptor {

    private static final Logger LOGGER = LoggerFactory.getLogger("sqlLog");

    private static final String TIME_SPLIT = "~";
    private static final AtomicLong COUNTER = new AtomicLong(0L);
    /** 每条 sql 执行前记录时间戳, 如果使用 ThreadLocal, 会有执行了 pre 但没有执行 post 的情况 */
    private static final Cache<Thread, String> TIME_CACHE = CacheBuilder.newBuilder().expireAfterWrite(60, TimeUnit.MINUTES).build();

    /** 多个空白符替换成一个空格 */
    private static final Pattern BLANK_REGEX = Pattern.compile("\\s{1,}");
    private static final String REPLACE = " ";

    @Override
    public void init(Connection connection, Properties properties) throws SQLException {}

    @Override
    public ResultSetInternalMethods preProcess(String sql, Statement statement, Connection connection) throws SQLException {
        if (LOGGER.isDebugEnabled()) {
            String realSql = getRealSql(sql, statement);
            if (isNotBlank(realSql)) {
                Thread currentThread = Thread.currentThread();
                long counter = COUNTER.addAndGet(1);
                long start = System.currentTimeMillis();

                TIME_CACHE.put(currentThread, counter + TIME_SPLIT + start);
                LOGGER.debug("计数: {}, db 主机: {}, sql: {}", counter, connection.getHost(), getLittleSql(realSql));
            }
        }
        return null;
    }
    private String getLittleSql(String sql) {
        int len = sql.length(), max = 2000, leftRight = 300;
        return len > max ? (sql.substring(0, leftRight) + " ... " + sql.substring(len - leftRight, len)) : sql;
    }

    @Override
    public ResultSetInternalMethods postProcess(String sql, Statement statement, ResultSetInternalMethods rs,
                                                Connection connection) throws SQLException {
        if (LOGGER.isDebugEnabled()) {
            String realSql = getRealSql(sql, statement);
            if (isNotBlank(realSql)) {
                Thread currentThread = Thread.currentThread();
                String counterAndTime = TIME_CACHE.getIfPresent(currentThread);
                if (isNotBlank(counterAndTime)) {
                    try {
                        String[] split = counterAndTime.split(TIME_SPLIT);
                        if (split.length == 2) {
                            long counter = toLong(split[0]);
                            long start = toLong(split[1]);

                            StringBuilder sbd = new StringBuilder();
                            if (counter > 0) {
                                sbd.append("计数: ").append(counter);
                            }
                            if (start > 0) {
                                sbd.append(", 时间: ").append(System.currentTimeMillis() - start).append("ms");
                            }
                            int size;
                            if (rs != null && rs.reallyResult() && rs.last()) {
                                size = rs.getRow();
                                rs.beforeFirst();

                                sbd.append(", 返回行数: ").append(size);
                            }
                            LOGGER.debug(sbd.toString());
                        }
                    } finally {
                        TIME_CACHE.invalidate(currentThread);
                    }
                }
            }
        }
        return null;
    }

    private String getRealSql(String sql, Statement statement) {
        if (isBlank(sql) && isNotBlank(statement)) {
            sql = statement.toString();
            if (isNotBlank(sql)) {
                int i = sql.indexOf(':');
                if (i > 0) {
                    sql = sql.substring(i + 1).trim();
                }
            }
        }
        if (isBlank(sql)) {
            return null;
        }

        // String realSql = SQLUtils.formatMySql(sql.replaceFirst("^\\s*?\n", ""));
        String realSql = BLANK_REGEX.matcher(sql.replaceFirst("^\\s*?\n", "")).replaceAll(REPLACE);
        return realSql.split("\n").length > 1 ? ("\n" + realSql) : realSql;
    }
    private boolean isBlank(Object obj) {
        return obj == null || obj.toString().trim().length() == 0;
    }
    private boolean isNotBlank(Object obj) {
        return !isBlank(obj);
    }
    private long toLong(Object obj) {
        if (isBlank(obj)) {
            return 0L;
        }

        try {
            return Long.parseLong(obj.toString().trim());
        } catch (NumberFormatException e) {
            return 0L;
        }
    }

    @Override
    public boolean executeTopLevelOnly() { return false; }

    @Override
    public void destroy() {
        TIME_CACHE.invalidateAll();
    }
}
